import sqlalchemy
import numpy as np

from .utils.model import train_model, predict_new_data
from .utils.read_db import generate_training_dataset, get_pre_lessons_grades
from .utils.read_db import get_stu_in_class, get_stu_info_dict

# 使用generate_training_dataset函数获取特征数据和目标变量
# 然后使用train_model函数训练和评估模型
# 最后使用predict_new_data函数预测新数据
def predict_for_stu_func(user_id: str, lesson_num: str, engine):
    X, y = generate_training_dataset(lesson_num, engine)
    X = np.array(X)
    y = np.array(y)
    model, img_path = train_model(X, y)
    pre_lesson_grades = get_pre_lessons_grades(lesson_num, engine)
    if user_id not in pre_lesson_grades:
        print(f"No grades found for user {user_id} in pre-lessons of {lesson_num}")
        return None
    user_grades = np.array([pre_lesson_grades[user_id]])
    predicted_grade = predict_new_data(model, user_grades)
    return predicted_grade[0], img_path


def predict_for_class_func(class_name: str, lesson_num: str, engine):
    # 获取班级中所有学生的user_id
    uid_in_class = get_stu_in_class(class_name, engine)
    # 获取特征数据和目标变量
    X, y = generate_training_dataset(lesson_num, engine)
    X = np.array(X)
    y = np.array(y)
    # 训练和评估模型
    try:
        model, img_path = train_model(X, y)
    except ValueError as e:
        print("not enough data to predict", e)
        return None
    # 获取班级学生的前置课程成绩
    pre_lesson_grades = get_pre_lessons_grades(lesson_num, engine)
    user_grades_list = []
    uid_has_grades = []
    # 准备预测数据
    for user_id in uid_in_class:
        if user_id in pre_lesson_grades:
            user_grades_list.append(pre_lesson_grades[user_id])
            uid_has_grades.append(user_id)
    # 如果没有学生有前置课程成绩，则返回空
    if len(user_grades_list) == 0:
        print(f"No grades found for any student in class {class_name} for pre-lessons of {lesson_num}")
        return None
    # 预测班级学生的成绩
    user_grades_array = np.array(user_grades_list)
    predicted_grades = predict_new_data(model, user_grades_array)
    # 将预测结果与用户ID对应起来
    predicted_grades_dict = {uid_has_grades[i]: predicted_grades[i] for i in range(len(uid_has_grades))}

    stu_info_dict = get_stu_info_dict(class_name, engine)
    result_list_with_info_list = []
    for uid, predicted_grade in predicted_grades_dict.items():
        stu_info = stu_info_dict.get(uid, {})
        result_list_with_info_list.append({
            "uid": uid,
            "predicted_grade": predicted_grade,
            "stu_num": stu_info.get("stu_num"),
            "username": stu_info.get("username")
        })

    return result_list_with_info_list, img_path
